import asyncio
import time

from lib import singleton
from .client import RpcClient
from typing import Dict, Optional
import random
from . import DiscoverOption, ServiceAddress, ServiceInfo
from urllib.parse import parse_qs

"""rpc调度器"""


@singleton
class RpcManage(object):
    def __init__(self, ip: str, port: int, discover_option: DiscoverOption, name=""):
        self.ip = ip
        self.port = port
        self.address = f"{self.ip}:{self.port}"
        self.name = name  # 作为谁的manage(worker？gateway？timer？)
        self.address2Client: Dict[str, RpcClient] = {}

        self.discover_node_name = discover_option.node_name
        self.discover_heartbeat_time = discover_option.heartbeat_time
        self.discover_instance = discover_option.instance
        self.watch_task: Optional[asyncio.Task] = None

    async def run(self):
        await self.discover()

    def close(self):
        self.watch_task.cancel()

    async def discover(self):
        """服务发现"""
        serviceInfo2Time = await self.discover_instance.get_all_service(self.discover_node_name)
        nt = int(time.time())
        for address, heartbeat_time in serviceInfo2Time.items():
            # 服务挂掉了？
            if nt - int(heartbeat_time) > self.discover_heartbeat_time:
                continue
            asyncio.create_task(self.connect_service(address))

        # 监听新的service
        self.watch_task = asyncio.create_task(
            self.discover_instance.watch(
                self.discover_node_name,
                callback=self.handle_service_info
            )
        )

    async def handle_service_info(self, service_info: ServiceInfo):
        service_info = {
            k: val[0]
            for k, val in parse_qs(service_info).items()
        }
        act, address = service_info["act"], service_info["address"]
        if act == "add":
            await self.connect_service(address)
        elif act == "del":
            self.del_service(address)

    async def connect_service(self, service_address: ServiceAddress):
        """连接到服务"""
        if service_address in self.address2Client:
            return

        ip, port = service_address.split(":")
        rpc_client = RpcClient(ip, int(port), tag=self.name, from_address=self.address)
        if await rpc_client.connect() == "fail":
            return

        self.address2Client[service_address] = rpc_client

    def del_service(self, service_address: ServiceAddress):
        """删除服务"""
        if service_address in self.address2Client:
            print(f"{self.name}[{self.address}] service {service_address} del")
            del self.address2Client[service_address]

    def rand_worker(self):
        """随机获取一个可用的client"""
        if not self.address2Client:
            return

        client_list = [c for c in self.address2Client.values() if c.working]
        if not client_list:
            return

        return random.choice(client_list)

    def rand_worker_url(self):
        """随机获取一个可用的worker url"""
        if not self.address2Client:
            return

        url_list = [url for url, c in self.address2Client.items() if c.working]
        if not url_list:
            return

        return random.choice(url_list)

    def get_worker(self, address):
        """获取指定address的client"""
        return self.address2Client.get(address)

    def is_ready(self, address):
        """address对应的服务是否可用"""
        return address in self.address2Client


if __name__ == '__main__':
    rpcManage = RpcManage("10.0.0.153", -1, name="test")


    async def start():
        await asyncio.gather(
            rpcManage.run(),
            asyncio.Future()
        )


    try:
        asyncio.run(start())
    except KeyboardInterrupt:
        pass
